const tf = require('@tensorflow/tfjs-node');
const mobilenet = require('@tensorflow-models/mobilenet'); // MobileNet 模型，用于迁移学习
const fs = require('fs-extra'); // 一个扩展了 Node.js 文件系统模块的库，用于文件操作

// 定义模型结构
const numClasses = 2;
const model = tf.sequential();
// model.add(tf.layers.flatten({ inputShape: [7, 7, 256] }));
model.add(tf.layers.flatten({ inputShape: [96, 96, 3] }));
model.add(tf.layers.dense({ units: 100, activation: 'relu' }));
model.add(tf.layers.dense({ units: numClasses, activation: 'softmax' }));

const labelArray = {
    cat: 0,
    dog: 1
}

// 编译模型
model.compile({
  optimizer: tf.train.adam(0.0001),
  loss: 'categoricalCrossentropy',
  metrics: ['accuracy'],
});
const trainDataPath = '../image/animal';
const testDataPath = '../image/animal-test';
const fileName = ['images']

let imagesTensor;
// 从文件夹中读取图片数据
async function getFileListAndLabel(path) {
    const imageData = []
    const files = await fs.readdirSync(path)
    const images = files.reduce((prev, dir) => {
        if (!fileName.includes(dir)) return
        let fileList = fs.readdirSync(path + '/' + dir)
        console.log(fileList)
        fileList.map(file => {
            // labels.push(dir)
            // return `${path}/${dir}/${file}`
            const n = file.lastIndexOf('.')
            if (n > -1) {
                const type = file.substring(n + 1)
                if (['jpg'].includes(type)) {
                    if (file[0] === file[0].toUpperCase()) {
                        imageData.push({
                            imagePath: `${path}/${dir}/${file}`,
                            label: 0,
                            className: file[0] === file[0].toUpperCase() ? 'cat' : 'dog'
                        })
                    } else {
                        imageData.unshift({
                            imagePath: `${path}/${dir}/${file}`,
                            label: 1,
                            className: file[0] === file[0].toUpperCase() ? 'cat' : 'dog'
                        })
                    }
                }
            }
        })
    }, [])
    return imageData
}
// 图片路径数据读取图片数据并转化为张量
const imageToTensor = (paths) => {
  // concat方法是将多个张量合并成一个张量
  return tf.concat(paths.reduce((prev, path) => {
    // 使用tidy方法，对张量进行内存控制，不然你直接读直接内存炸
    try {
        const tensor = tf.tidy(() => {
          return tf.node
          .decodeImage(fs.readFileSync(path.imagePath)) // 把图片转换成张量
          .resizeNearestNeighbor([96, 96]) // 把图片的宽高转换成96x96
          .toFloat() // 转换成浮点型的张量
          .div(tf.scalar(255.0)) // 将图片张量进行归一化（这个自行去理解）
          .expandDims()
        })
        if (tensor.shape.length !== 4 || tensor.shape[3] !== 3) {
            console.log('tensor', path, tensor.shape)
        }
        prev.push(tensor)
        return prev
    } catch {
        console.log('path', path.imagePath)
    }
  }, []))
}
// 从文件中加载训练数据
async function loadTrainData(dirPath) {
  try {
    const data = await getFileListAndLabel(dirPath)
    // const labelList = data.map(item => item.label)
    // console.log('labelList', labelList.join('-'))
    imagesTensor = imageToTensor(data)
    console.log('imagesTensor', imagesTensor)

    // 将标签转换为 TensorFlow.js 的 Tensor
    const labelList = data.map(item => item.label)
    const labels = tf.tensor1d(labelList, 'int32');
    const classNames = data.map(item => item.className);
    // 创建 tf.data.Dataset
    // const dataset = tf.data.zip({ images: images, labels: labels });
    const dataset = { images: imagesTensor, labels };
    console.log('dataset', dataset)
    // 返回 Dataset
    return dataset
  } catch (error) {
    console.error('加载训练数据时出错:', error);
    throw error;
  }
}
// loadTrainData(trainDataPath)

async function start() {
    // 准备数据
    const trainData = await loadTrainData(trainDataPath)/* 从文件中加载训练数据 */;
    const testData = await loadTrainData(testDataPath)/* 从文件中加载测试数据 */;
    // console.log(trainData.labels)
    // 将标签转换为独热编码
    const trainLabels = tf.oneHot(trainData.labels, numClasses);
    const testLabels = tf.oneHot(testData.labels, numClasses);
    
    // 训练模型
    model.fit(trainData.images, trainLabels, {
      epochs: 15,
      batchSize: 32,
      validationData: [testData.images, testLabels],
      shuffle: true, // 打乱数据
    }).then(info => {
      console.log('训练完成:', info);
      // 保存模型
      model.save('file://model').then(() => {
        console.log('模型保存成功。');
      });
    });
}
start()